
#include "bitwise_left_shift_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

const int BLOCK_SIZE = 32;
const int MAX_DIM = 5;
const int MAX_CORES = 40;
const int PRESERVED = 8 * 1024;

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context) {
    // 用于控制UB buffer的使用，默认为不需要广播的情况
    int TENSOR_NUM = 3;
    int BUFFER_NUM = 1;

    const gert::StorageShape* input_shape = context->GetInputShape(0);
    const gert::StorageShape* other_shape = context->GetInputShape(1);

    uint32_t nDimInput = input_shape->GetStorageShape().GetDimNum();
    uint32_t nDimOther = other_shape->GetStorageShape().GetDimNum();
    uint64_t ub_size;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);

    //初始化shape信息
    uint32_t shapeInput[MAX_DIM] = {};
    uint32_t shapeOther[MAX_DIM] = {};
    for (int i = 0; i < MAX_DIM; i++) {
        shapeInput[i] = 1;
        shapeOther[i] = 1;
    }

    //通过数组传递shape信息到npu侧
    uint32_t typeSize = GetSizeByDataType(context->GetInputDesc(0)->GetDataType());
    uint64_t totalDataSize = typeSize;
    for (int i = 0; i < nDimInput; i++) {
        shapeInput[i] = input_shape->GetStorageShape().GetDim(i);
        totalDataSize *= shapeInput[i];
    }
    uint32_t offset = nDimInput - nDimOther;
    for (int i = 0; i < nDimOther; i++) {
        shapeOther[i + offset] = other_shape->GetStorageShape().GetDim(i);
    }

    //确认数据是否需要广播
    int needBroadcast = 0;
    int tilingKey = 0;
    int nBatch, batchDataSize;
    int broadcastDim = 0;
    nBatch = batchDataSize = 1;
    for (int i = 0; i < nDimInput; i++) {
        if (shapeInput[i] != shapeOther[i]) { //需要广播
            needBroadcast++;
        }
    }
    if (needBroadcast == 1) { // 只需要广播一个维度
        tilingKey = 3;
        int i;
        for (i = 0; i < nDimInput; i++) {
            nBatch *= shapeInput[i];
            if (shapeInput[i] != shapeOther[i]) {
                broadcastDim = shapeInput[i];
                break;
            }
        }
        for (int j = i + 1; j < nDimInput; j++) {
            batchDataSize *= shapeInput[j];
        }
    } else if (needBroadcast > 1) { //需要广播的维度大于1
        tilingKey = 1;
        TENSOR_NUM = 2;
        BUFFER_NUM = 1;
    } else if (!needBroadcast && (typeSize == 2 || typeSize == 4)) { //不需要广播，对于int16和int32，可采用向量化计算
        tilingKey = 2;
        // bitmask和临时位移结果需要的空间，尽可能保留较大的空间，因为Compare需要256byte对齐，会开更大的buffer
        TENSOR_NUM = 7;
    } else { // 不需要广播，无法向量化计算
        tilingKey = 0;
    }
    context->SetTilingKey(tilingKey);

    //处理除单维度广播以外的所有情况下的数据划分，直接按照输入数据展开至一维进行划分
    uint16_t maxBlockPerIter = (ub_size - PRESERVED) / BUFFER_NUM / TENSOR_NUM / BLOCK_SIZE;
    // 对数据根据input进行划分，与广播无关
    uint32_t nCores = MAX_CORES;
    uint32_t totalBlocks = (totalDataSize + BLOCK_SIZE - 1) / BLOCK_SIZE;
    while (totalBlocks <= 5 * nCores && nCores > 1)
        nCores /= 2;
    uint32_t blockPerCore = totalBlocks / nCores;
    uint16_t nAcores = totalBlocks % nCores;
    uint16_t nBcores = nCores - nAcores;

    // 设置tiling
    if (tilingKey == 2 || tilingKey == 0) {
        BitwiseLeftShiftTilingData tiling;
        tiling.set_blockPerCore(totalBlocks / nCores);
        tiling.set_maxBlockPerIter(maxBlockPerIter);
        tiling.set_nAcores(nAcores);
        tiling.set_nBcores(nBcores);
        tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
        context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    } else if (tilingKey == 3) {
        BitwiseLeftShiftTilingDataBroadSpec tiling;
        if (nBatch < nCores)
            nCores = nBatch; //如果batch数量小于核数，则使用batch数量作为核数
        nAcores = nBatch % nCores;
        nBcores = nCores - nAcores;
        tiling.set_broadcastDim(broadcastDim);
        tiling.set_batchDataSize(batchDataSize);
        tiling.set_nBatchPerCore(nBatch / nCores);
        tiling.set_nAcores(nAcores);
        tiling.set_nBcores(nBcores);
        tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
        context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    } else {
        BitwiseLeftShiftTilingDataBroad tiling;
        tiling.set_blockPerCore(totalBlocks / nCores);
        tiling.set_maxBlockPerIter(maxBlockPerIter);
        tiling.set_nAcores(nAcores);
        tiling.set_nBcores(nBcores);
        tiling.set_nDim(nDimInput);
        tiling.set_shapeInput(shapeInput);
        tiling.set_shapeOther(shapeOther);
        tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
        context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    }
    //设置核数
    context->SetBlockDim(nCores);
    return ge::GRAPH_SUCCESS;
}
} // namespace optiling

namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context) {
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
static ge::graphStatus InferDataType(gert::InferDataTypeContext* context) {
    const auto inputDataType = context->GetInputDataType(0);
    context->SetOutputDataType(0, inputDataType);
    return ge::GRAPH_SUCCESS;
}
} // namespace ge

namespace ops {
class BitwiseLeftShift : public OpDef {
  public:
    explicit BitwiseLeftShift(const char* name) : OpDef(name) {
        this->Input("input")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT8, ge::DT_INT16, ge::DT_INT32, ge::DT_INT64})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("other")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT8, ge::DT_INT16, ge::DT_INT32, ge::DT_INT64})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("out")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT8, ge::DT_INT16, ge::DT_INT32, ge::DT_INT64})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);

        this->AICore().SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");
    }
};

OP_ADD(BitwiseLeftShift);
} // namespace ops
